# 北京炒家吃首板模式
import matplotlib.pyplot as plt
import akshare as ak
from datacache import get_stock_data

def plot_results(df, portfolio):
    """
    绘制回测结果
    :param df: 股票数据
    :param portfolio: 回测结果
    """
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
    plt.figure(figsize=(14, 7))
    plt.subplot(2, 1, 1)
    plt.plot(df['close'], label='Close Price', color='blue')
    plt.title('Stock Close Price Over Time')
    plt.xlabel('Date')
    plt.ylabel('Price')
    plt.legend()
    plt.grid(True)

    # 绘制收益曲线（累计收益率）
    plt.subplot(2, 1, 2)
    plt.plot(df['sum'], label='Cumulative Return', color='green')
    plt.title('Cumulative Returns Over Time')
    plt.xlabel('Date')
    plt.ylabel('Cumulative Return')
    plt.legend()
    plt.grid(True)

    plt.tight_layout()
    plt.show()

def main():
    # stock_sse_summary_df = ak.stock_sse_summary()
    # print(stock_sse_summary_df)
    # 参数设置
    stock_code = 'sz002276'  # 平安银行
    start_date = '20220901'
    end_date = '20250101'

    # 获取数据
    df = get_stock_data(stock_code, start_date, end_date)
    df['sum'] = 0
    sum = [] # 收益曲线
    initial_capital = 10000 # 初始资金
    holdings = 0
    is_prepare_sell = False # 上一个交易日是涨停
    print(df.columns.to_list())
    for index, row in df.iterrows():
        row_date = row['date']
        close_price = row['close']
        open_price = row['open']
        high_price = row['high'] # 最高价
        low_price = row['low'] # 最低价
        volume = row['volume']
        turnover = row['turnover'] # 换手率

        volume_rate = 1
        if index > 5:
            value1 = df['volume'].values[index - 5]
            value2 = df['volume'].values[index - 4]
            value3 = df['volume'].values[index - 3]
            value4 = df['volume'].values[index - 2]
            value5 = df['volume'].values[index - 1]
            average_volumn = (value1 + value2 + value3 + value4 + value5) / 6
            volume_rate = volume / average_volumn

        # 缩量上涨还要涨，所以不卖
        if is_prepare_sell and volume_rate > 1:
            is_prepare_sell = False
            initial_capital = initial_capital + holdings * close_price
            holdings = 0
            print("卖出于:" + row_date)

        # 持有仓位并且收盘价大于最高价，买入，并且下一个交易日卖出 考虑缩量持续上涨 并且上涨超过5个点才操作（稳妥） 3%风险高收益高
        elif holdings == 0 and is_prepare_sell == False and volume_rate > 1 and (close_price - open_price) / open_price > 0.05 and (close_price - open_price) / open_price < 0.095:
            is_prepare_sell = True
            holdings = initial_capital / close_price
            initial_capital = initial_capital - holdings * close_price
            print("卖入于:" + row_date)
        df.loc[index, 'sum'] = initial_capital + holdings * close_price
        # sum.append(initial_capital + holdings * close_price)
        sum.append(volume)
    plot_results(df, sum)

if __name__ == "__main__":
    main()